Skip to content

feat: pass calling SessionContext to Python UDTF callbacks#1555

Draft
timsaucer wants to merge 1 commit into
apache:mainfrom
timsaucer:feat/df54-followups-wave3
Draft

feat: pass calling SessionContext to Python UDTF callbacks#1555
timsaucer wants to merge 1 commit into
apache:mainfrom
timsaucer:feat/df54-followups-wave3

Conversation

@timsaucer
Copy link
Copy Markdown
Member

Which issue does this PR close?

Related to #1533

Rationale for this change

Upstream DataFusion deprecated call() on table functions so that we could support far more features. Instead we should use call_with_args. That is already being done in main using default table function arguments. This PR exposes the full functionality of the function signature.

What changes are included in this PR?

  • Enhance table functions to take the table parameters.
  • Add unit tests

Are there any user-facing changes?

Yes, this changes the expected arguments to pure python table functions.

DataFusion 53 added `TableFunctionImpl::call_with_args(TableFunctionArgs)`
where `TableFunctionArgs` carries both the positional expression
arguments and the calling `&dyn Session`. The pure-Python UDTF path
previously discarded everything but the exprs.

Thread the session through when the user callback's signature opts in
by declaring a `session` keyword parameter (or `**kwargs`). At call
time we downcast the `&dyn Session` to its canonical `SessionState`
impl and build a fresh `SessionContext` over the same Arc-shared state,
exposed to Python as a `datafusion.SessionContext` wrapper. Existing
callbacks whose signatures do not declare `session` continue to be
called with the positional expression arguments only — no behavior
change for current users.

Note: a UDTF body cannot drive a fresh `ctx.sql(...).collect()` on the
passed-in session because the outer SQL execution already holds the
tokio runtime. Use the session for metadata access (catalogs, UDF
lookups, config) rather than nested DataFrame collection.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant